import argparse
import json
import os
import shutil
import subprocess
import sys
from glob import glob

# 默认 ffmpeg / ffprobe 路径（你提供的）
FFMPEG_DIR = r"C:\Soft\ffmpeg-2025-09-25-git-9970dc32bf-full_build"
FFMPEG = os.path.join(FFMPEG_DIR, "bin", "ffmpeg.exe")
FFPROBE = os.path.join(FFMPEG_DIR, "bin", "ffprobe.exe")

VIDEO_EXTS = {".mp4", ".mkv", ".mov", ".avi", ".m4v", ".webm", ".ts", ".flv", ".mpg", ".mpeg", ".mkv", ".wmv"}

def run(cmd):
    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    return p.returncode, p.stdout, p.stderr

def probe_streams(path):
    # 增加帧率信息用于计算第一帧时长
    cmd = [FFPROBE, "-v", "error", "-show_entries", "stream=index,codec_type,codec_name,disposition,r_frame_rate,avg_frame_rate", "-of", "json", path]
    rc, out, err = run(cmd)
    if rc != 0:
        raise RuntimeError(f"ffprobe failed: {err.strip()}")
    info = json.loads(out)
    return info.get("streams", [])

def _parse_rate(rate_str):
    try:
        if not rate_str:
            return 0.0
        if isinstance(rate_str, (int, float)):
            return float(rate_str)
        if "/" in rate_str:
            num, den = rate_str.split("/", 1)
            num = float(num) if num else 0.0
            den = float(den) if den else 1.0
            return num / den if den != 0 else 0.0
        return float(rate_str)
    except Exception:
        return 0.0

def get_video_frame_rate(streams):
    for s in streams:
        if s.get("codec_type") == "video":
            # 优先使用 r_frame_rate，再用 avg_frame_rate
            fps = _parse_rate(s.get("r_frame_rate") or s.get("avg_frame_rate"))
            return fps
    return 0.0

def build_map_indices(streams):
    keep = []
    for s in streams:
        disp = s.get("disposition", {})
        if disp.get("attached_pic", 0) == 1:
            # skip attached pic
            continue
        keep.append(s["index"])
    return keep

def remove_first_frame(src, dst, streams=None):
    try:
        if streams is None:
            streams = probe_streams(src)
    except Exception as e:
        return ("error", f"ffprobe failed when probing for frame rate: {e}")

    fps = get_video_frame_rate(streams)
    if fps <= 0:
        # 兜底默认 25fps
        fps = 25.0
    frame_time = 1.0 / fps
    ss_str = f"{frame_time:.6f}"

    # 尝试无损 remux（-ss 放在 -i 之后，尽量只跳过开头）
    cmd = [FFMPEG, "-y", "-i", src, "-ss", ss_str, "-c", "copy", dst]
    rc, out, err = run(cmd)
    if rc == 0:
        return ("ok", f"Removed first frame (remux): {os.path.basename(src)}")

    # 无损失败，回退到快速编码（重编码视频，拷贝音频）
    cmd2 = [FFMPEG, "-y", "-ss", ss_str, "-i", src, "-c:v", "libx264", "-preset", "veryfast", "-crf", "23", "-c:a", "copy", dst]
    rc2, out2, err2 = run(cmd2)
    if rc2 == 0:
        return ("ok", f"Removed first frame (re-encoded video): {os.path.basename(src)}")
    return ("error", err2.strip() or err.strip() or out2.strip() or out.strip())

def remux_without_cover(src, dst):
    streams = probe_streams(src)
    keep = build_map_indices(streams)
    if len(keep) == len(streams):
        # 没有 attached_pic —— 按要求删除第一帧
        return remove_first_frame(src, dst, streams=streams)
    if not keep:
        return ("no_streams", f"No streams left after removing covers for {os.path.basename(src)}")
    cmd = [FFMPEG, "-y", "-i", src]
    for idx in keep:
        cmd += ["-map", f"0:{idx}"]
    cmd += ["-c", "copy", dst]
    rc, out, err = run(cmd)
    if rc != 0:
        return ("error", err.strip() or out.strip())
    return ("ok", "")

def find_files(paths):
    files = []
    for p in paths:
        if os.path.isdir(p):
            for root, _, filenames in os.walk(p):
                for fn in filenames:
                    if os.path.splitext(fn)[1].lower() in VIDEO_EXTS:
                        files.append(os.path.join(root, fn))
        else:
            # support glob patterns
            for match in glob(p):
                if os.path.isfile(match):
                    files.append(match)
    return sorted(set(files))

def choose_output_dir(base_name="nocover", parent_dir=None):
    if parent_dir is None:
        parent_dir = os.getcwd()
    idx = 0
    while True:
        name = base_name if idx == 0 else f"{base_name}{idx}"
        path = os.path.join(parent_dir, name)
        if not os.path.exists(path):
            try:
                os.makedirs(path, exist_ok=True)
            except Exception as e:
                raise RuntimeError(f"cannot create output directory {path}: {e}")
            return path
        idx += 1

def main():
    parser = argparse.ArgumentParser(description="Batch remove attached picture (cover) using ffmpeg (lossless remux).")
    parser.add_argument("paths", nargs="+", help="Files, globs or directories to process")
    parser.add_argument("--overwrite", "-o", action="store_true", help="Overwrite original files (safe temp then replace)")
    parser.add_argument("--ffmpeg", help="Path to ffmpeg executable (optional)")
    parser.add_argument("--ffprobe", help="Path to ffprobe executable (optional)")
    parser.add_argument("--outdir", help="Parent directory to create nocover folder(s) in (default: C:\\Users\\zyc\\Videos)")
    args = parser.parse_args()

    global FFMPEG, FFPROBE
    if args.ffmpeg:
        FFMPEG = args.ffmpeg
    if args.ffprobe:
        FFPROBE = args.ffprobe

    # 默认 parent dir 改为 C:\Users\zyc\Videos，支持通过 --outdir 覆盖
    parent_for_output = args.outdir if args.outdir else r"C:\Users\zyc\Videos"
    try:
        output_dir = choose_output_dir("nocover", parent_dir=parent_for_output)
    except Exception as e:
        print(f"Failed to prepare output directory: {e}")
        return 2
    print(f"Output directory: {output_dir}")

    files = find_files(args.paths)
    if not files:
        print("No files found.")
        return 1

    for f in files:
        base = os.path.basename(f)  # 保持原文件名不变
        outpath = os.path.join(output_dir, base)
        # 如果输出文件已经存在则覆盖（保持行为简单），或者你可以改为添加后缀
        status, msg = remux_without_cover(f, outpath)
        if status == "ok":
            if args.overwrite:
                # 用生成的文件替换原文件（谨慎）
                tmp = outpath + ".tmp_replace"
                try:
                    shutil.move(outpath, tmp)
                    os.replace(tmp, f)
                    print(f"[OK] Overwrote original: {base}")
                except Exception as e:
                    print(f"[ERR] Replacing failed for {base}: {e}")
            else:
                print(f"[OK] Created in {output_dir}: {base} - {msg}")
        elif status == "no_streams":
            print(f"[SKIP] {msg}")
            if os.path.exists(outpath):
                os.remove(outpath)
        else:
            print(f"[ERR] {base}: {msg}")
            if os.path.exists(outpath):
                os.remove(outpath)
    return 0

if __name__ == "__main__":
    sys.exit(main())